Restructure valid waveform mask in Heterodyne likelihood by SSL32081 · Pull Request #125 · GW-JAX-Team/Jim

SSL32081 · 2026-07-10T19:32:19Z

This PR refactors much of the code of the masking applied to the waveform and frequencies during initialisation of the Heterodyned Likelihood. This mask aimed to mask out non-physical/invalid and trivial waveform pieces, which was introduced in kazewong/jim@93023114b.

The mask is defined as:

h_amp = jnp.sum(
    jnp.array([jnp.abs(h_sky[pol]) for pol in h_sky.keys()]), axis=0
)
f_valid = frequency_original[jnp.where(h_amp > 0)[0]]

mask_heterodyne_center = jnp.where(
    (self.freq_grid_center <= jnp.max(f_valid))
    & (self.freq_grid_center >= jnp.min(f_valid))
)[0]

At the time, it was necessary to mask out the NaNs from the waveform, since the self.frequencies attribute at the time begins at zero (see e.g. here), where the waveform evaluation is invalid. But this is no longer needed with the introduction of sliced_frequencies in kazewong/jim@b745e04.

That being said, the mask is still needed to remove the waveform where it is trivial (= 0), which occurs soon after the merger frequency. For otherwise it would lead to NaN (from 0/0) later when taking the ratio between waveform pieces.

Summary by CodeRabbit [edited]

Improvements
- Enhanced consistency and stability of frequency binning for heterodyned transient likelihood calculations.
- Improved likelihood coefficient normalisation for more robust evaluation across detector duration.
Refactor
- Updated bin-boundary handling to simplify and align coefficient preparation while preserving expected results.
Tests
- Strengthened likelihood initialisation checks (including finiteness validation) and updated test expectations for low-frequency reference cutoff behaviour.

for more information, see https://pre-commit.ci

coderabbitai · 2026-07-10T19:32:41Z

Important

Review skipped

Review was skipped due to path filters

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

CodeRabbit blocks several paths by default. You can override this behavior by explicitly including those paths in the path filters. For example, including **/dist/** will override the default block on the dist directory, by removing the pattern from both the lists.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 6069a113-df94-497f-a867-39901af7fadb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

📝 Walkthrough

Walkthrough

HeterodynedTransientLikelihoodFD now uses masked bin edges as the primary binning representation, derives bin centres internally, and normalises coefficients using detector duration. Tests validate finite stored arrays and the updated coefficient interface.

Changes

Heterodyne bin-edge representation

Layer / File(s)	Summary
Bin-edge binning initialisation `src/jimgw/core/single_event/likelihood.py`	Likelihood annotations and initialisation now derive valid frequency bounds, widths, reference responses, and coefficients from masked bin edges.
Binning helper and coefficient contracts `src/jimgw/core/single_event/likelihood.py`	`_make_binning_scheme` returns edges only; `_mask_and_set_frequency_arrays` establishes valid frequency metadata; `_compute_coefficients` derives centres internally and uses detector duration for normalisation.
Likelihood validation updates `tests/unit/core/single_event/test_likelihood.py`	Tests check finite stored detector arrays and adapt the coefficient test double to receive bin edges.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

GW-JAX-Team/Jim#116: Reworks related relative-binning and heterodyne logic in HeterodynedTransientLikelihoodFD.

Suggested reviewers: narolaharsh

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: restructuring the valid waveform mask in the Heterodyne likelihood.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch rm-obsolete-mask

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/jimgw/core/single_event/likelihood.py`:
- Around line 697-699: In HeterodynedTransientLikelihoodFD.__init__, reorder the
assignments so self.freq_grid_low and self.freq_grid_high are set from freq_grid
before calculating self.bin_widths; then derive bin_widths from those
initialized attributes.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 9b56adbe-eb3a-4fa0-a369-4536edb125be

📥 Commits

Reviewing files that changed from the base of the PR and between 4d7f611 and 4cf526c.

📒 Files selected for processing (1)

src/jimgw/core/single_event/likelihood.py

for more information, see https://pre-commit.ci

coderabbitai

Actionable comments posted: 3

🧹 Nitpick comments (1)

src/jimgw/core/single_event/likelihood.py (1)
556-561: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Remove the spaces from these jaxtyping shape strings — " n_bins + 1" / " n_valid+1" are parsed as separate axes, so these annotations need compact symbolic expressions. Ruff already ignores F722, so no config change is needed.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/jimgw/core/single_event/likelihood.py` around lines 556 - 561, The
jaxtyping shape annotations in the affected likelihood definitions use spaces
around symbolic expressions, causing them to parse as separate axes. Update the
relevant annotations, including `freq_grid_low`, `freq_grid_high`, `bin_widths`,
`waveform_low_ref`, `waveform_high_ref`, and `summary_data`, to use compact
shape strings such as `"n_bins+1"` and `"n_valid+1"` without spaces; leave Ruff
configuration unchanged.
Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/jimgw/core/single_event/likelihood.py`:
- Around line 849-875: Remove the self references from _compute_coefficients,
using the supplied f_bins to derive the bin edges and centers needed for the
mask and frequency-shift calculations. Replace self.freq_grid_low and
self.freq_grid_high with values computed from f_bins while preserving the
final-bin inclusive equality handling.
- Line 817: In the relevant likelihood class initialization logic, replace the
assignment to self.nbins with self.n_bins so the valid-bin count is stored in
the declared attribute expected by downstream callers and tests.

In `@tests/unit/core/single_event/test_likelihood.py`:
- Around line 1174-1234: Restore and activate
test_low_frequency_reference_cutoff_does_not_reindex_summary_data, updating its
assertions and setup to match the final HeterodynedTransientLikelihoodFD
contract; retain coverage for masked bin-edge reindexing, reduced n_bins,
frequency bounds, and per-detector coefficient dimensions, and remove the
commented-out test block.

---

Nitpick comments:
In `@src/jimgw/core/single_event/likelihood.py`:
- Around line 556-561: The jaxtyping shape annotations in the affected
likelihood definitions use spaces around symbolic expressions, causing them to
parse as separate axes. Update the relevant annotations, including
`freq_grid_low`, `freq_grid_high`, `bin_widths`, `waveform_low_ref`,
`waveform_high_ref`, and `summary_data`, to use compact shape strings such as
`"n_bins+1"` and `"n_valid+1"` without spaces; leave Ruff configuration
unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ee9ecd0d-aa26-4e37-800c-91668c5e69b5

📥 Commits

Reviewing files that changed from the base of the PR and between 4cf526c and 2d34459.

📒 Files selected for processing (2)

src/jimgw/core/single_event/likelihood.py
tests/unit/core/single_event/test_likelihood.py

SSL32081 and others added 2 commits July 11, 2026 03:16

Remove redundant waveform mask

322b143

[pre-commit.ci] auto fixes from pre-commit.com hooks

4cf526c

for more information, see https://pre-commit.ci

SSL32081 changed the base branch from main to jim-dev July 10, 2026 19:32

SSL32081 requested a review from thomasckng July 10, 2026 19:34

SSL32081 self-assigned this Jul 10, 2026

SSL32081 added the enhancement New feature or request label Jul 10, 2026

coderabbitai Bot reviewed Jul 10, 2026

View reviewed changes

Comment thread src/jimgw/core/single_event/likelihood.py Outdated

SSL32081 and others added 10 commits July 11, 2026 03:46

Remove redundant variable definitions

32cfc2d

Fix freq_grid reference before assignment

80468f3

Passing detector object directly in compute_coefficients

cdd2b74

[pre-commit.ci] auto fixes from pre-commit.com hooks

869d357

for more information, see https://pre-commit.ci

Update likelihood test

b88ec69

Adding back frequency mask

d46118d

[pre-commit.ci] auto fixes from pre-commit.com hooks

2d34459

for more information, see https://pre-commit.ci

Fix merge conflict

a730282

Minor order change

f840eaf

Restore indices test in likelihood test

f2fd27c

coderabbitai Bot reviewed Jul 11, 2026

View reviewed changes

Comment thread src/jimgw/core/single_event/likelihood.py Outdated

Comment thread src/jimgw/core/single_event/likelihood.py

Comment thread tests/unit/core/single_event/test_likelihood.py Outdated

SSL32081 added 2 commits July 11, 2026 15:38

Fix self reference in staticmethod

d9816c5

Rename nbins to n_bins

a0c0f1e

SSL32081 changed the title ~~Remove obsolete valid waveform mask from Heterodyne likelihood~~ Restructure valid waveform mask in Heterodyne likelihood Jul 11, 2026

SSL32081 and others added 4 commits July 11, 2026 15:43

Minor syntax fixes

270bd08

Fix freq_bins in compute coefficients

1f22d81

Update lockfile

6ebfcfd

Merge branch 'jim-dev' into rm-obsolete-mask

cefddbe

thomasckng merged commit ead70b6 into jim-dev Jul 13, 2026
3 of 6 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Restructure valid waveform mask in Heterodyne likelihood#125

Restructure valid waveform mask in Heterodyne likelihood#125
thomasckng merged 18 commits into
jim-devfrom
rm-obsolete-mask

SSL32081 commented Jul 10, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 10, 2026 •

edited

Loading

Review skipped

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

SSL32081 commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit [edited]

Uh oh!

coderabbitai Bot commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

SSL32081 commented Jul 10, 2026 •

edited

Loading

coderabbitai Bot commented Jul 10, 2026 •

edited

Loading